Skip to content

feat: unified chunk grid with rectilinear chunk/shard support#3802

Open
maxrjones wants to merge 148 commits intozarr-developers:mainfrom
maxrjones:poc/unified-chunk-grid
Open

feat: unified chunk grid with rectilinear chunk/shard support#3802
maxrjones wants to merge 148 commits intozarr-developers:mainfrom
maxrjones:poc/unified-chunk-grid

Conversation

@maxrjones
Copy link
Copy Markdown
Member

@maxrjones maxrjones commented Mar 21, 2026

Summary

This PR contains an alternative implementation of the rectilinear chunk grid extension, building on the work in #3534 (RLE helpers, validation logic, and test cases were directly adopted). While the core feature of variable-sized chunks is the same, the internal architecture differs in ways that impact extensibility, performance, and release safety.

I appreciate the patience of those who contributed to #3534, and everyone who's been waiting on this feature. I know it's frustrating to see a new PR after #3534 was so close. That PR provided fundamental components, and I hope people will see the value here. I really believe it is worth the churn for the following reasons:

Key differences from #3534

  1. Extensibility. Each dimension is represented by a type implementing the DimensionGrid protocol (FixedDimension, VaryingDimension). Adding a new dimension type (e.g. TiledDimension for periodic patterns like days-per-month) requires implementing that protocol — no changes to indexing, codecs, or the ChunkGrid class. A prototype was built to verify this.
  2. Performance. The indexing pipeline queries each dimension independently with scalar calls rather than constructing N-d coordinate tuples per chunk lookup. This avoids allocation overhead in the inner loop of every indexer. VaryingDimension uses precomputed prefix sums for O(log n) lookups via binary search. See https://github.com/maxrjones/zarr-chunk-grid-tests for a performance comparison.
  3. Feature flag. Rectilinear chunk grids are gated behind zarr.config.set({'array.rectilinear_chunks': True}) (or ZARR_ARRAY__RECTILINEAR_CHUNKS=True), disabled by default. This gives downstream libraries time to adapt before the API is finalized, and us an opportunity to gracefully finalize the API.
  4. Rectilinear sharding. Shard boundaries can be rectilinear while inner chunks remain regular, with validation that each shard edge is divisible by the inner chunk size. This is tested end-to-end and documented in the user guide.

Design document: docs/design/chunk-grid.md covers the full design, rationale, and a suggested PR sequence for splitting this into reviewable increments, if needed.

Downstream POCs (all passing):

TODO:

  • Add unit tests and/or doctests in docstrings
  • Add docstrings and API docs for any new/modified user-facing classes and functions
  • New/modified features documented in docs/user-guide/*.md
  • Changes documented as a new file in changes/
  • GitHub Actions have all passed
  • Test coverage is 100% (Codecov passes)

@maxrjones
Copy link
Copy Markdown
Member Author

thank you very much for the review @dcherian! I believe that I have addressed each comment via either a change or a todo/follow-up issue. Please let me know if you see anything else or have remaining concerns.

@codecov
Copy link
Copy Markdown

codecov bot commented Apr 3, 2026

Codecov Report

❌ Patch coverage is 94.35262% with 41 lines in your changes missing coverage. Please review.
✅ Project coverage is 92.98%. Comparing base (ca9bd3e) to head (dbd2cf8).

Files with missing lines Patch % Lines
src/zarr/testing/stateful.py 20.00% 12 Missing ⚠️
src/zarr/core/metadata/v3.py 92.75% 10 Missing ⚠️
src/zarr/core/chunk_grids.py 97.42% 6 Missing ⚠️
src/zarr/core/array.py 95.41% 5 Missing ⚠️
src/zarr/core/common.py 90.90% 4 Missing ⚠️
src/zarr/core/metadata/v2.py 25.00% 3 Missing ⚠️
src/zarr/codecs/sharding.py 90.90% 1 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #3802      +/-   ##
==========================================
+ Coverage   92.91%   92.98%   +0.06%     
==========================================
  Files          86       87       +1     
  Lines       10715    11246     +531     
==========================================
+ Hits         9956    10457     +501     
- Misses        759      789      +30     
Files with missing lines Coverage Δ
src/zarr/abc/codec.py 98.76% <ø> (ø)
src/zarr/api/synchronous.py 92.95% <ø> (ø)
src/zarr/codecs/transpose.py 91.52% <ø> (ø)
src/zarr/core/_info.py 97.59% <100.00%> (+2.40%) ⬆️
src/zarr/core/codec_pipeline.py 94.18% <ø> (ø)
src/zarr/core/config.py 100.00% <ø> (ø)
src/zarr/core/group.py 94.85% <ø> (ø)
src/zarr/core/indexing.py 96.12% <100.00%> (+0.04%) ⬆️
src/zarr/experimental/__init__.py 100.00% <100.00%> (ø)
src/zarr/metadata/migrate_v3.py 98.36% <100.00%> (ø)
... and 8 more
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@d-v-b
Copy link
Copy Markdown
Contributor

d-v-b commented Apr 6, 2026

@maxrjones from my POV you can merge this whenever you are ready. I believe this does technically contain breaking changes to the codec pipeline API (the signature of validate has changed, and necessarily IMO), but I don't think any downstream packages were actually using this parameter, so I believe this is safe to go in.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

7 participants